import json
import os
import shutil
import time

from PyQt6.QtCore import QThreadPool

import midjourney_config
from global_config import get_workflows_img_dir, get_workflows_svd_dir
from scripts import db_utils
from scripts.WorkSingle import WorkerSignal, Worker
from scripts.configui import comfyui_client
from scripts.configui.apis import comfyui_api
from scripts.gpts import Gpt_Utils


class GenerateManager:

    def __init__(self, callback):
        self.callback = callback
        self.signal = WorkerSignal()
        self.signal.finished.connect(self.mergeimagelist)
        self.tasklist = list[absTask]()

    def on_message(self, ws, message):
        # print(f"Received message: {message}")
        jsonobj = json.loads(message)
        type = jsonobj['type']
        data = jsonobj['data']

        if type == 'execution_start':
            1
            # print(jsonobj)
        elif type == 'executing':
            if data['node'] is None:
                self.signal.finished.emit(data['prompt_id'])
        elif type == 'progress':
            1
            # print(data['value'] / data['max'])

    # 定义出错时的回调函数
    def on_error(self, ws, error):
        ws.close()
        print(f"Error: {error}")

    # 启动任务
    def startTask(self):
        comfyui_client.start_comfyui_client(self.on_message, self.on_error)

    def enQueue(self, task):
        task.dotask()
        self.tasklist.append(task)
        # worker = Worker(inner=lambda: (), callback=lambda a: ())
        # thread = QThreadPool.globalInstance()
        # thread.start(worker)

    # 将换脸后的图片 与 不用换脸的图片 ，按顺序合并到list
    def mergeimagelist(self, prompt_id):
        for i, task in enumerate(self.tasklist):
            if task.prompt_id == prompt_id and not task.isDone:
                task.process(self.callback, len(self.tasklist) == 1)
                break

        self.tasklist = [task for task in self.tasklist if not task.isDone]


class absTask:
    isDone: bool = False
    prompt_id = None

    def dotask(self):
        pass

    def process(self, callback, isEnd):
        pass


class ImageTask(absTask):

    def __init__(self, prompts, negative, index, photoname=None, width=None, height=None):
        self.taskid = index
        self.photoName = photoname
        self.prompts = prompts
        self.negative = negative
        self.width = width
        self.height = height

    def dotask(self):
        project_config = db_utils.load_project_config()
        type_wash = project_config['type']
        type = project_config['reasoning_platform']
        userid = db_utils.load_project_config()['userid']
        # 洗稿模式，不翻译
        content = self.prompts
        if midjourney_config.type_origin == type_wash :
            conversation_id, content = Gpt_Utils.translate_prompts(type, userid, userid, self.prompts)
        if self.photoName is None:
            self.prompt_id = comfyui_api.text2img_workflow(content, self.negative, self.width, self.height)  # 文生图
        else:
            self.prompt_id = comfyui_api.img2img_workflow(self.photoName, content, self.negative, self.width,
                                                              self.height)
        time.sleep(0.2)

    def process(self, callback, isEnd):
        print('callback ======', self.taskid, self.prompt_id)
        # 获取Comfyui 生成的图片list
        result_images = comfyui_api.get_history(self.prompt_id)
        # print(result_images)
        # 复制到项目路径
        imagedir = get_workflows_img_dir()
        taskdir = f'{imagedir}\\{self.taskid}\\'
        if not os.path.exists(taskdir):
            os.mkdir(taskdir)

        newlist = []
        for img in result_images:
            image = comfyui_api.get_view_image(img)

            newpath = os.path.join(taskdir, img)
            image.save(newpath)
            newlist.append(newpath)

        callback(True, newlist, taskdir, self.taskid, isEnd, False)
        self.isDone = True


class SVDTask(absTask):
    def __init__(self, taskid, imagePath, index):
        self.taskid = taskid
        self.imagePath = imagePath
        self.index = index

    def dotask(self):
        self.prompt_id = comfyui_api.svd_workflow(self.imagePath, self.index)
        time.sleep(0.2)

    def process(self, callback, isEnd):

        # 获取Comfyui 生成的图片list
        result_images = comfyui_api.get_history(self.prompt_id)
        # 复制到项目路径
        svddir = get_workflows_svd_dir()
        taskdir = f'{svddir}\\{self.taskid}\\'
        if not os.path.exists(taskdir):
            os.mkdir(taskdir)

        newlist = []
        for img in result_images:
            content = comfyui_api.get_view_video(img)
            if content is not None:
                newpath = os.path.join(taskdir, img)
                with open(newpath, 'wb') as f:
                    f.write(content)
                newlist.append(newpath)

        callback(True, newlist, taskdir, self.taskid, isEnd, True)
        self.isDone = True
